import requests
import json

RM_IP = "10.1.212.193"
API_PORT = "8088"


def get_yarn_applications():
    url = f"http://{RM_IP}:{API_PORT}/ws/v1/cluster/apps"
    try:
        resp = requests.get(url)
        resp.raise_for_status()
        apps = resp.json().get("apps", {}).get("app", [])

        # 计算任务等待时间（仅PENDING任务）
        pending_durations = [
            app["elapsedTime"] for app in apps
            if app["state"] in ("NEW", "NEW_SAVING")
        ]
        max_pending_time = max(pending_durations) if pending_durations else 0

        # 按状态分类统计
        app_stats = {
            "所有任务数": len(apps),
            "运行中任务数": len([app for app in apps if app["state"] == "RUNNING"]),
            "失败任务数": len([app for app in apps if app["state"] == "FAILED"]),
            "手动终止任务数": len([app for app in apps if app["state"] == "KILLED"]),
            "等待资源任务数": len([app for app in apps if app["state"] == "NEW" or app["state"] == "NEW_SAVING"]),
            "最大等待秒数": max_pending_time // 1000,
        }
        return app_stats
    except Exception as e:
        print(f"获取任务信息失败: {e}")
        return {}

def get_yarn_metrics():
    url = f"http://{RM_IP}:{API_PORT}/ws/v1/cluster/metrics"
    try:
        resp = requests.get(url, timeout=5)
        resp.raise_for_status()  # 检查HTTP状态码
        metrics = resp.json()
    except requests.exceptions.RequestException as e:
        print(f"API请求失败: {e}")
        return {}

    totalMB = metrics["clusterMetrics"]["totalMB"]
    totalVirtualCores = metrics["clusterMetrics"]["totalVirtualCores"]
    utilizedMBPercent = metrics["clusterMetrics"]["utilizedMBPercent"]
    utilizedVirtualCoresPercent = metrics["clusterMetrics"]["utilizedVirtualCoresPercent"]
    totalNodes = metrics["clusterMetrics"]["totalNodes"]
    activeNodes = metrics["clusterMetrics"]["activeNodes"]
    lostNodes = metrics["clusterMetrics"]["lostNodes"]
    unhealthyNodes = metrics["clusterMetrics"]["unhealthyNodes"]
    allocatedMB=metrics["clusterMetrics"]["allocatedMB"]
    allocatedVirtualCores=metrics["clusterMetrics"]["allocatedVirtualCores"]

    yarn_dict = {
        "总内存GB": round(totalMB / 1024, 2),
        "已用内存GB": round(allocatedMB / 1024, 2),
        "内存使用率": utilizedMBPercent,
        "总CPU": totalVirtualCores,
        "已用CPU": allocatedVirtualCores,
        "CPU使用率": utilizedVirtualCoresPercent,
        "节点数": totalNodes,
        "活动节点数": activeNodes,
        "丢失节点数": lostNodes,
        "异常节点数": unhealthyNodes,
        **get_yarn_applications()  # 合并任务统计
    }
    print(yarn_dict)


get_yarn_metrics()
